package cn.apienvrypt.advice;

import java.util.Objects;
import javax.annotation.Resource;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.core.MethodParameter;
import org.springframework.http.MediaType;
import org.springframework.http.converter.HttpMessageConverter;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.web.bind.annotation.ControllerAdvice;
import org.springframework.web.servlet.mvc.method.annotation.ResponseBodyAdvice;
import com.google.gson.Gson;

import cn.apienvrypt.annotation.Encrypt;
import cn.apienvrypt.config.SecretKeyConfig;
import cn.apienvrypt.context.HandleContext;
import cn.apienvrypt.domain.ResultDto;
import cn.apienvrypt.util.AESUtil;
import cn.apienvrypt.util.GsonUtil;
import cn.apienvrypt.util.RSAUtil;

/**
 * @author lfj
 */
@ControllerAdvice
public class EncryptResponseBodyAdvice implements ResponseBodyAdvice<Object> {

	private final Logger log = LoggerFactory.getLogger(this.getClass());
	private boolean encrypt;
	@Resource
	private SecretKeyConfig secretKeyConfig;

	@Override
	public boolean supports(MethodParameter returnType, Class<? extends HttpMessageConverter<?>> converterType) {
		encrypt = Objects.requireNonNull(returnType.getMethod())
			.isAnnotationPresent(Encrypt.class) && secretKeyConfig.isOpen();
		return encrypt;
	}

	@Override
	public Object beforeBodyWrite(Object body, MethodParameter returnType, MediaType selectedContentType,
		Class<? extends HttpMessageConverter<?>> selectedConverterType, ServerHttpRequest request, ServerHttpResponse response) {

		if(encrypt) {
			try {
				if(!(body instanceof ResultDto)) {
					throw new RuntimeException("无法识别的响应体,请使用ResultDto构造返回");
				}

				Gson gson = GsonUtil.getGson();
				ResultDto<?> resultDtoSource = (ResultDto<?>) body;
				String resultDtoMsg = gson.toJson(resultDtoSource.getData());

				// 获取前端请求中的aesKey加密data
				String aesKey = HandleContext.currentValue();
				String encryptDataStr = AESUtil.getInstance()
					.encrypt(aesKey, resultDtoMsg);

				// 使用RSA私钥签名,前端验签
				String signMsg = RSAUtil.sign(encryptDataStr, RSAUtil.getPrivateKey(secretKeyConfig.getPrivateKey()));

				ResultDto<String> resultDto = new ResultDto<>();
				resultDto.setData(encryptDataStr);
				resultDto.setMsgCode(resultDtoSource.getMsgCode());
				resultDto.setPage(resultDtoSource.isPage());
				resultDto.setSign(signMsg);

				return resultDto;
			} catch(Exception e) {
				log.error("加密响应数据异常:", e);
			} finally {
				HandleContext.remove();
			}
		}

		return body;
	}
}
